from PIL import Image
import pillow_avif
import argparse
import os
import sys
import glob
import uuid  # 添加uuid模块用于生成唯一哈希值

# 导入颜色工具模块
from color_utils import RED, GREEN, YELLOW, BLUE, PURPLE, CYAN, BOLD, NC

def generate_unique_hash():
    """生成一个唯一的哈希值"""
    return uuid.uuid4().hex[:8]  # 使用8位哈希值，可以根据需要调整长度

def convert_image_format(input_path, output_path=None, target_format='png'):
    """
    转换图片格式
    
    Args:
        input_path: 输入图片路径
        output_path: 输出图片路径，如果为None则自动生成
        target_format: 目标格式，如'png', 'jpg', 'webp', 'avif', 'ico', 'tif'等
    
    Returns:
        bool: 处理是否成功
    """
    try:
        # 打开输入图片
        input_image = Image.open(input_path)
        
        # 如果没有指定输出路径，则根据输入文件名和哈希值生成
        if output_path is None:
            input_filename = os.path.basename(input_path)
            name, _ = os.path.splitext(input_filename)
            unique_hash = generate_unique_hash()
            output_path = os.path.join(os.path.dirname(input_path), f"{name}_{unique_hash}.{target_format.lower()}")
        
        # 确保输出目录存在
        output_dir = os.path.dirname(output_path)
        if output_dir and not os.path.exists(output_dir):
            os.makedirs(output_dir)
        
        # 如果指定了输出路径，则直接使用；否则添加唯一哈希值
        if output_path:
            new_output_path = output_path
        else:
            filename, ext = os.path.splitext(output_path)
            unique_hash = generate_unique_hash()
            new_output_path = f"{filename}_{unique_hash}{ext}"
        
        # 如果是JPEG格式，移除透明通道
        if target_format.lower() in ['jpg', 'jpeg']:
            if input_image.mode in ['RGBA', 'LA']:
                # 创建白色背景
                background = Image.new('RGB', input_image.size, (255, 255, 255))
                if input_image.mode == 'RGBA':
                    background.paste(input_image, mask=input_image.split()[3])
                else:
                    background.paste(input_image, mask=input_image.split()[1])
                input_image = background
        
        # 特殊格式处理
        if target_format.lower() == 'ico':
            # ICO格式需要特殊处理，确保尺寸合适
            sizes = [(16, 16), (32, 32), (48, 48), (64, 64)]
            if input_image.size[0] > 256 or input_image.size[1] > 256:
                # 如果图片太大，先缩小
                input_image.thumbnail((256, 256), Image.LANCZOS)
            input_image.save(new_output_path, format='ICO', sizes=sizes)
        elif target_format.lower() == 'jpg':
            input_image.convert('RGB').save(new_output_path, 'JPEG', quality=95)
        elif target_format.lower() == 'tif' or target_format.lower() == 'tiff':
            input_image.save(new_output_path, 'TIFF', compression='tiff_lzw')
        elif target_format.lower() == 'avif':
            input_image.save(new_output_path, 'AVIF', quality=80)
        else:
            input_image.save(new_output_path, target_format.upper())
        
        print(f"{GREEN}格式转换成功: {BLUE}{os.path.basename(input_path)}{NC} -> {BLUE}{os.path.basename(new_output_path)}{NC}")
        return True
    except Exception as e:
        print(f"{RED}处理图片时出错: {str(e)}{NC}")
        return False

def batch_convert(input_dir, output_dir, target_format='png'):
    """
    批量转换目录下的所有图片格式
    
    Args:
        input_dir: 输入目录路径
        output_dir: 输出目录路径
        target_format: 目标格式
        
    Returns:
        bool: 是否全部处理成功
    """
    # 支持的图片格式
    image_extensions = ['*.jpg', '*.jpeg', '*.png', '*.bmp', '*.gif', '*.webp', '*.tiff', '*.tif', '*.avif', '*.ico', '*.heic']
    
    # 获取所有图片文件
    image_files = []
    for ext in image_extensions:
        image_files.extend(glob.glob(os.path.join(input_dir, ext)))
        image_files.extend(glob.glob(os.path.join(input_dir, ext.upper())))
    
    if not image_files:
        print(f"{RED}错误: 在目录 {BLUE}{input_dir}{NC} 中未找到任何图片文件{NC}")
        return False
    
    print(f"{CYAN}找到 {BOLD}{len(image_files)}{NC}{CYAN} 个图片文件，开始批量转换格式...{NC}")
    
    # 确保输出目录存在
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    # 处理每个图片
    success_count = 0
    for input_file in image_files:
        try:
            # 获取文件名
            filename = os.path.basename(input_file)
            name, _ = os.path.splitext(filename)
            
            # 设置输出文件路径，添加唯一哈希值
            unique_hash = generate_unique_hash()
            output_file = os.path.join(output_dir, f"{name}_{unique_hash}.{target_format.lower()}")
            
            print(f"{CYAN}转换: {BLUE}{filename}{NC} -> {BLUE}{os.path.basename(output_file)}{NC}")
            
            # 转换格式
            if convert_image_format(input_file, output_file, target_format):
                success_count += 1
        except Exception as e:
            print(f"{RED}处理 {BLUE}{input_file}{NC} 时出错: {str(e)}{NC}")
    
    print(f"{GREEN}批量转换完成: 成功 {BOLD}{success_count}/{len(image_files)}{NC}{GREEN} 个文件{NC}")
    return success_count > 0

def print_help():
    """打印帮助信息"""
    print(f"{CYAN}{BOLD}图片格式转换工具 - 转换图片格式{NC}")
    print("")
    print(f"{BOLD}用法:{NC}")
    print(f"  {YELLOW}xzimg format -i <输入文件> -o <输出文件> [-f <格式>]{NC}")
    print(f"  {YELLOW}xzimg format -b [输入目录] -o [输出目录] [-f <格式>]{NC}")
    print("")
    print(f"{BOLD}参数:{NC}")
    print(f"  {GREEN}-i, --input{NC} <文件>    指定输入图片文件路径")
    print(f"  {GREEN}-o, --output{NC} <文件>   指定输出图片文件路径或目录")
    print(f"  {GREEN}-b, --batch{NC} [目录]    批量处理模式，可选指定输入目录，默认为当前目录")
    print(f"  {GREEN}-f, --format{NC} <格式>   指定目标格式，默认为png")
    print(f"  {GREEN}--help{NC}               显示此帮助信息")
    print("")
    print(f"{BOLD}支持的格式:{NC}")
    print(f"  {PURPLE}png, jpg/jpeg, webp, gif, bmp, tiff/tif, avif, ico, heic{NC}")
    print("")
    print(f"{BOLD}示例:{NC}")
    print(f"  {YELLOW}xzimg format -i input.png -o output.jpg -f jpg{NC}")
    print(f"  {YELLOW}xzimg format -b -f webp{NC}                       # 将当前目录下的所有图片转换为WebP格式")
    print(f"  {YELLOW}xzimg format -b images -o converted -f jpg{NC}    # 将images目录下的所有图片转换为JPG格式")

def main():
    # 创建参数解析器
    parser = argparse.ArgumentParser(description='图片格式转换工具 - 转换图片格式', add_help=False)
    
    # 添加帮助参数
    parser.add_argument('--help', action='store_true', help='显示帮助信息')
    
    # 创建互斥组，确保-i和-b不会同时使用
    group = parser.add_mutually_exclusive_group()
    group.add_argument('-i', '--input', help='输入图片路径')
    group.add_argument('-b', '--batch', nargs='?', const=os.getcwd(), help='批量处理的输入目录路径，不指定值则使用当前目录')
    
    parser.add_argument('-o', '--output', help='输出图片路径或目录')
    parser.add_argument('-f', '--format', default='png', help='目标格式，如png, jpg, webp等，默认为png')
    
    # 解析命令行参数
    args, unknown = parser.parse_known_args()
    
    # 如果有--help参数，显示帮助信息并退出
    if args.help:
        print_help()
        return 0
    
    # 检查必要参数
    if not args.input and not args.batch:
        print(f"{RED}错误: 必须指定输入文件(-i)或批量处理模式(-b){NC}")
        print(f"使用 {YELLOW}--help{NC} 查看帮助信息")
        return 1
    
    # 判断是单文件模式还是批量模式
    if args.input:
        # 单文件模式
        if not os.path.isfile(args.input):
            print(f"{RED}错误: 输入文件不存在: {BLUE}{args.input}{NC}")
            return 1
        
        # 处理图片 - 如果没有指定输出路径，则传入None让函数自动生成
        success = convert_image_format(args.input, args.output, args.format)
    else:
        # 批量模式
        input_dir = args.batch
        
        # 如果没有指定输出目录，默认为当前目录下的format文件夹
        if not args.output:
            output_dir = os.path.join(os.getcwd(), 'format')
        else:
            output_dir = args.output
        
        if not os.path.isdir(input_dir):
            print(f"{RED}错误: 输入目录不存在: {BLUE}{input_dir}{NC}")
            return 1
        
        # 批量处理图片
        success = batch_convert(input_dir, output_dir, args.format)
    
    # 返回状态码
    return 0 if success else 1

if __name__ == "__main__":
    sys.exit(main())